tomosphero
TomoSphero
A 3D/4D volume raytracer in spherical coordinates for arbitrary detector shape implemented in PyTorch.
Check the tutorial for instruction on using this library or examples for complete samples demonstrating forward raytracing and retrieval.
Features
- 3D spherical raytracing with optional support for dynamic volume (4D)
- implemented purely in PyTorch for easy integration with PyTorch's optimization and machine learning capabilities
- support for square/circular detectors or other custom detector shapes
- optional retrieval framework for easily defining loss functions and parametric models (currently supports only static 3D volumes)
Quickstart
pip install tomosphero
git clone https://github.com/evidlo/tomosphero && cd tomosphero
python examples/single_vantage.py
python examples/static_retrieval.py
python examples/dynamic_measurements.py
Memory Usage
This library was uses only PyTorch array operations for implementation simplicity and speed at the expense of memory consumption. The peak memory usage in GB can be approximated with examples/memory_usage.py
$ python examples/memory_usage.py
--- Parameters ---
(50, 50, 50) object
50 observations, 1 channels, (50, 100) sensor
--- Memory Usage ---
Ray coordinates memory: 4.25 GB
Object memory: 0.05 GB
Architecture
Below is a list of modules in this package and their purpose:
Forward Raytracing
raytracer.py- computation of voxel indices for intersecting rays, raytracing Operatorgeometry.py- viewing geometry (detector) definitions, spherical grid definition
Retrieval
model.py- parameterized models for representing a density. used in retrievalloss.py- some loss functions to be used in retrievalretrieval.py- retrieval algorithmsplotting.py- functions for plotting stacks of images, retrieval losses
Running Tests
pytest tomosphero
See Also
tomosipo, which inspired parts of this library's interface.
Tutorial
This tutorial will walk through the process of setting up a tomography problem and performing a simple static reconstruction. See the examples/ directory if you want complete working scripts.
The tutorial consists of four parts:
- Construct a spherical grid that defines the object domain
- Construct view geometry which defines measurement LOS
- Combine grid and view geometry into an operator and take measurements
- Reconstruct object with iterative algorithm
Let's get started! 🮲🮳
Constructing Spherical Grids
The grid defines the physical extent and shape of the object to be raytraced or reconstructed. There are several ways to instantiate grids with the SphericalGrid class in TomoSphero:
By shape and size. e.g. for a spherical grid with 30 bins in radius/elevation/azimuth and outer radius of 10 (arbitrary) length units.
from tomosphero import SphericalGrid
grid = SphericalGrid(
# voxels (radius, elevation, azimuth)
shape=(30, 30, 30),
size_r=(0, 1), # inner/outer grid radius
# size_e=(0, 2*pi), # (radians, default)
# size_a=(-pi, pi), # (radians, default)
)
grid.plot()
Grids do not need to cover the full extent of a sphere. Below is an example of a grid defined over a small spherical wedge. This can be useful for reconstructions on a local region of a planetary atmosphere atmosphere.
from numpy import pi
from tomosphero import SphericalGrid
grid = SphericalGrid(
# voxels (radius, elevation, azimuth)
shape=(10, 10, 10),
size_r=(3, 10), # inner/outer grid radius
size_e=(0.4 * pi, 0.6 * pi), # (radians)
size_a=(-.1 * pi, .1 * pi), # (radians)
)
grid.plot()
Alternatively, grids may be specified by defining the specific location of boundaries in radius/elevation/azimuth dimensions.
from numpy import pi
from tomosphero import SphericalGrid
grid = SphericalGrid(
r_b=[8, 9, 10],
e_b=[0.4*pi, 0.5*pi, 0.6*pi],
a_b=[-0.1*pi, 0.0*pi, 0.1*pi],
)
grid.plot()
Grid shapes may be queried later with grid.shape and grid.size. See tomosphero.grid.SphericalGrid.
Constructing View Geometries
A view geometry defines the lines of sight associated with each pixel in a sensor. TomoSphero currently has 4 built-in view geometry types (contributions welcome):
tomosphero.geometry.ConeRectGeom- Conventional rectangular camera sensor- lines of sight converge to single point at detector location
tomosphero.geometry.ParallelGeom- Parallel lines of sight. Common geometry for medical imagingtomosphero.geometry.ConeCircGeom- Exotic circular virtual camera geometry used by the Carruthers spacecrafttomosphero.geometry.ViewGeom- Base class for specifying arbitrary LOS. Sensor geometry may have any dimension/shape
ConeRectGeom is a rectilinear camera sensor, the most common sensor type. It is defined by giving sensor shape, location, orientation, and field of view:
from tomosphero import ConeRectGeom
geom = ConeRectGeom(
shape=(64, 64), # pixels
pos=(10, 0, 0), # sensor location
# lookdir=(-1, 0, 0) # sensor orientation (default: towards origin)
fov=(45, 45) # degrees
)
geom.plot()
View geometries may be composed by addition to produce scanning orbits (e.g. circular/helical orbits), The result of this operation is a tomosphero.geometry.ViewGeomCollection object, which behaves much like a regular ViewGeom.
import numpy as np
from tomosphero import ConeRectGeom
geom = None
# sweep 360° around origin in orbit of radius 5
for w in np.linspace(0, 2*np.pi, 50):
pos=(5*np.cos(w), 5*np.sin(w), 2)
# combine geoms by addition
geom += ConeRectGeom(
pos=pos,
shape=(100, 100), # pixels
fov=(25, 25) # degrees
)
anim = geom.plot()
View geometry shape may be queried later with geom.shape, and its internal LOS are available through geom.ray_starts for LOS start points and geom.rays for normal vectors.
Taking Projections
The tomosphero.raytracer.Operator class is responsible for carrying out the raytracing operation, taking the grid and view geometry specified previously as arguments.
At instantiation, this object will compute and store the intersection coordinates of all lines of sight with all grid boundaries and may take some time depending on the sensor and grid sizes.
TomoSphero makes extensive use of PyTorch for its autograd capabilities and GPU support. TomoSphero uses the CPU by default, but the GPU may be selected with the device kwarg.
Let's first begin by creating an Operator and defining a 3D object in a PyTorch tensor.
For simplicity, let's take our object to be a sphere with a wedge missing from it:
# compute LOS intersections over grid
op = Operator(grid, geom, device='cuda')
# example phantom object - broken torus
rho = t.ones(grid.shape, device='cuda')
# sideways pacman object
rho[:, :, 3:] = 1
After instantiation, the operator accepts PyTorch arrays with shape grid.shape (and with appropriate device) and returns raytraced measurements.
# raytrace and get measurements
y = op(rho)
# y.shape matches geom.shape
# plot measurements
from tomosphero.plotting import image_stack
anim = image_stack(y, geom)
Reconstruction
TomoSphero is designed to be used as a component in iterative tomographic reconstruction and provides some basic building blocks for setting up an inverse problem. Given the synthetically-generated measurement $y$ from the last section, we can use this iterative approach to fit a reconstruction $\hat{\rho}$ to our measurements $y$.
Compared to more conventional techniques such as filtered back-projection, (S)ART-based methods, an iterative reconstruction approach built on an autograd framework has many benefits:
- Works with any view geometry
- Filtered back-projection techniques are generally limited to circular or helical orbits
- Availability of different optimizers
- Any PyTorch optimizer] may be used interchangeably in TomoSphero. Some optimizers may handle certain problems better than others when it comes to e.g. nonconvexity.
- Rapid testing of different cost functions and regularizers
- (S)ART-based methods are hard coded to a specific loss function
- Iterative autograd techniques can work with any differentiable loss/regularizers
We begin by instantiating a tensor with requires_grad=True so PyTorch will track its gradients and set up the optimizer:
rho_hat = t.zeros(grid.shape, requires_grad=True)
optim = t.optim.Adam([rho_hat], lr=1e-1)
We will use a very basic least-squares loss function with no model parameterization or regularizers for our reconstruction:
$$\hat{\rho} = \arg \min_{\rho} ||y - F \rho||_2^2$$
where $y$ are measurements, $\rho$ is the reconstructed object, and $F$ is the tomographic operator.
for i in range(100):
optim.zero_grad()
# compute loss and its gradient
loss = t.sum((y - op(rho_hat))**2)
loss.backward()
optim.step()
anim = image_stack(preview3d(rho))
plt.title('Ground Truth')
anim = image_stack(preview3d(rho_hat))
plt.title('Reconstruction')
We have successfully reconstructed the object! Of course, we had ideal measurement conditions with good angular coverage and no noise, but these cases can be handled by the introduction of more sophisticated models and regularization.
As mentioned above, one of the major benefits of autograd frameworks is the flexibility to change any part of the loss function or model with minimal changes to the code. Model-based retrievals provides more details on TomoSphero's tools that make experimenting with parametric models and regularizers easier.
Miscellaneous
Retrieval Framework
TomoSphero provides an optional object-oriented framework for conveniently experimenting with new loss functions and tracking loss history in iterative retrieval. The image below shows an example of a plot created with the loss tracking tools for a problem with three loss terms:
Plotting Utilities
TomoSphero features extensive plotting capabilities with many objects containing a built-in .plot() function for preview. In addition, TomoSphero has a tomosphero.plotting module with several more plotting utilities:
tomosphero.plotting.image_stack- animate a 3D stack of 2D imagestomosphero.plotting.color_negative- convert grayscale positive/negative images to red/green RGB imagestomosphero.plotting.preview3d- generate rotating preview of an object and grid